import os
"""
该脚本文件必须放在训练集的根目录下，即与data.yaml文件同级
"""

"""
在当前目录下生成指定txt文件，文件内容是递归遍历所有子目录下的所有图片文件名，并写入txt文件;

root_dir: 根目录（train的上级目录）
output_txt: 输出的路径文件
"""
def generate_txt(root_dir, output_txt):
    with open(output_txt, "w") as f:
        # 遍历所有子目录
        for batch_dir in os.listdir(root_dir):
            batch_path = os.path.join(root_dir, batch_dir)
            if os.path.isdir(batch_path):
                # 图像目录
                img_dir = os.path.join(batch_path, "images")
                if os.path.exists(img_dir):
                    # 遍历图像文件
                    for img_name in os.listdir(img_dir):
                        if img_name.endswith((".jpg", ".jpeg", ".png")):
                            # 生成相对路径（相对于data.yaml）
                            img_path = os.path.join(img_dir, img_name)
                            # 写入txt
                            f.write(f"{img_path}\n")

    print(f"已生成路径文件：{output_txt}，包含{len(open(output_txt).readlines())}张图像")

if __name__ == "__main__":
    generate_txt("./train", "./train.txt")
    generate_txt("./val", "./val.txt")